# encoding=utf-8
import re
import nltk
from nltk.corpus import stopwords
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB

# 定义数据集
pos_reviews = [
    "I love this product, it's amazing!",
    "This is the best product I have ever used!",
    "I highly recommend this product.",
    "This product is awesome.",
    "This product is great and worth every penny.",
]
neg_reviews = [
    "I hate this product, it's terrible!",
    "This is the worst product I have ever used!",
    "I do not recommend this product.",
    "This product is terrible and not worth the money.",
    "This product is a waste of money.",
]

# 将数据集拼接为一个字符串列表，并标记情感类别
all_reviews = []
for review in pos_reviews:
    all_reviews.append((review, 'positive'))
for review in neg_reviews:
    all_reviews.append((review, 'negative'))

# 对文本进行预处理，去除停用词，提取词干
stopwords = set(stopwords.words('english'))
stemmer = nltk.PorterStemmer()


def preprocess(text):
    text = re.sub(r'[^\w\s]', '', text)  # 去除标点符号
    words = nltk.word_tokenize(text.lower())  # 分词并转为小写
    words = [stemmer.stem(word) for word in words if word not in stopwords]  # 提取词干并去除停用词
    return ' '.join(words)


# 对所有文本进行预处理，并转换为文本向量
corpus = [preprocess(review[0]) for review in all_reviews]
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(corpus).toarray()

# 定义训练集和测试集
train_size = int(len(all_reviews) * 0.8)
train_X = X[:train_size]
train_y = [review[1] for review in all_reviews][:train_size]
test_X = X[train_size:]
test_y = [review[1] for review in all_reviews][train_size:]

# 训练朴素贝叶斯分类器
classifier = MultinomialNB()
classifier.fit(train_X, train_y)

# 在测试集上进行预测
predictions = classifier.predict(test_X)

# 计算准确率和召回率
tp = 0  # 正确预测为正样本的数量
tn = 0  # 正确预测为负样本的数量
fp = 0  # 错误预测为正样本的数量
fn = 0  # 错误预测为负样本的数量
for i in range(len(predictions)):
    print(predictions[i])
    if predictions[i] == 'positive' and test_y[i] == 'positive':
        tp += 1
    elif predictions[i] == 'negative' and test_y[i] == 'negative':
        tn += 1
    elif predictions[i] == 'positive' and test_y[i] == 'negative':
        fp += 1
    elif predictions[i] == 'negative' and test_y[i] == 'positive':
        fn += 1
# precision = tp / (tp + fp)
# recall = tp / (tp + fn)
# print('Precision:', precision)
# print('Recall:', recall)
print('tp:', tp)
print('fp:', fp)
print('fn:', fn)
